#include "CodecNvJPEG2000.h"

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "nvjpeg2k.h"

#ifndef QT_VERSION
#include <iostream>
#endif
#include <float.h>

class JP2_U16C1Encoder::D
{
public:
    bool isvalid = true;
    int cudaDevNum = 0;
    cudaStream_t cudastream = nullptr;
    nvjpeg2kEncoder_t enc_handle = nullptr;
    nvjpeg2kEncodeState_t enc_state = nullptr;
    nvjpeg2kEncodeParams_t enc_params = nullptr;
    double psnr = 0;        //  压缩前后图像峰值信号噪声比，  越小压缩比越大
    void *input_bufDev[1] = {nullptr};  //  可以为多通道，  这个类为单通道
    size_t input_bufDevPitch[1] = {0};
    nvjpeg2kImageComponentInfo_t image_comp_info;
    nvjpeg2kImage_t input_image;
    nvjpeg2kEncodeConfig_t enc_config;
    size_t output_buf_size = 1;  //  编码后的输出缓冲区大小
    unsigned char * output_buf = nullptr;  //  编码后的输出缓冲区
    size_t encoded_size = 0;    //  编码后的实际大小
};

JP2_U16C1Encoder::JP2_U16C1Encoder(int width,int height,int cudaDevNum):
    m_dptr(new D)
{
    if(width<1) width = 1;
    if(height<1) height = 1;

    std::memset(&m_dptr->image_comp_info, 0, sizeof(nvjpeg2kImageComponentInfo_t));
    std::memset(&m_dptr->input_image, 0, sizeof(nvjpeg2kImage_t));
    std::memset(&m_dptr->enc_config, 0, sizeof(nvjpeg2kEncodeConfig_t));

    m_dptr->cudaDevNum = cudaDevNum;
    int count = 0;
    int driverVersion = 0;
    cudaError_t ret = cudaGetDeviceCount(& count);
    cudaDriverGetVersion(&driverVersion);
    m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;
    m_dptr->isvalid = m_dptr->isvalid && count > 0;
    m_dptr->isvalid = m_dptr->isvalid && m_dptr->cudaDevNum < count;
    m_dptr->isvalid = m_dptr->isvalid && driverVersion>=1100;

    if(m_dptr->isvalid){
        ret = cudaSetDevice(m_dptr->cudaDevNum);
        m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;

        auto j2k_ret = nvjpeg2kEncoderCreateSimple(&m_dptr->enc_handle);   //  编码器名柄
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;

        j2k_ret = nvjpeg2kEncodeStateCreate(m_dptr->enc_handle,&m_dptr->enc_state);  //  编码器状态
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;

        j2k_ret = nvjpeg2kEncodeParamsCreate(&m_dptr->enc_params);    //  编码器
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;

        ret = cudaStreamCreate(&m_dptr->cudastream);
        m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;

        ret = cudaMallocPitch(&m_dptr->input_bufDev[0], &m_dptr->input_bufDevPitch[0],
                static_cast<size_t>(width * 2),static_cast<size_t>(height)
                );
//        m_dptr->input_bufDevPitch[0] = static_cast<size_t>(width * 2);
//        ret = cudaMalloc(&m_dptr->input_bufDev[0],static_cast<size_t>(width * 2 * height));

        m_dptr->isvalid = m_dptr->isvalid && ret==cudaSuccess;

        m_dptr->image_comp_info.component_width = static_cast<unsigned int>(width);
        m_dptr->image_comp_info.component_height = static_cast<unsigned int>(height);
        m_dptr->image_comp_info.precision = 16;
        m_dptr->image_comp_info.sgn = 0;

        m_dptr->input_image.pixel_data = m_dptr->input_bufDev;
        m_dptr->input_image.pixel_type = NVJPEG2K_UINT16;
        m_dptr->input_image.pitch_in_bytes = m_dptr->input_bufDevPitch;
        m_dptr->input_image.num_components = 1;   // 单通道

        m_dptr->enc_config.stream_type      =  NVJPEG2K_STREAM_JP2; // the bitstream will be in JP2 container format
        m_dptr->enc_config.color_space      =  NVJPEG2K_COLORSPACE_GRAY; // input image is in gray format
        m_dptr->enc_config.image_width      =  m_dptr->image_comp_info.component_width;
        m_dptr->enc_config.image_height     =  m_dptr->image_comp_info.component_height;
//        m_dptr->enc_config.enable_tiling = 1;
//        m_dptr->enc_config.tile_width = m_dptr->image_comp_info.component_width;
//        m_dptr->enc_config.tile_height = m_dptr->image_comp_info.component_height;
        m_dptr->enc_config.num_layers = 1;
        m_dptr->enc_config.num_components   =  1;
        m_dptr->enc_config.image_comp_info  =  &m_dptr->image_comp_info;
        m_dptr->enc_config.code_block_w     =  32;   //  设置成 32 x 32 编码速度可以快2倍以上
        m_dptr->enc_config.code_block_h     =  32;
        m_dptr->enc_config.irreversible     =  0;   //  小波变换是否不可逆,Reversible(5-3 lossless) and irreversible(9-7 lossy) wavelet transforms
        m_dptr->enc_config.mct_mode         =
                m_dptr->enc_config.color_space == NVJPEG2K_COLORSPACE_SRGB ? 1 : 0;
        m_dptr->enc_config.prog_order       =  NVJPEG2K_LRCP;
        m_dptr->enc_config.num_resolutions  =  6;

        j2k_ret = nvjpeg2kEncodeParamsSetEncodeConfig(m_dptr->enc_params,
                                                      &m_dptr->enc_config);
        m_dptr->isvalid = m_dptr->isvalid && NVJPEG2K_STATUS_SUCCESS == j2k_ret;

        m_dptr->output_buf_size = static_cast<size_t>(width * height * 2 * 2);
        m_dptr->output_buf = new unsigned char[m_dptr->output_buf_size];
    }
}

JP2_U16C1Encoder::~JP2_U16C1Encoder()
{
    if(m_dptr->isvalid){
        cudaStreamDestroy(m_dptr->cudastream);
    }

    if(m_dptr->enc_params){
        nvjpeg2kEncodeParamsDestroy(m_dptr->enc_params);
        m_dptr->enc_params = nullptr;
    }
    if(m_dptr->enc_state){
        nvjpeg2kEncodeStateDestroy(m_dptr->enc_state);
        m_dptr->enc_state = nullptr;
    }
    if(m_dptr->enc_handle){
        nvjpeg2kEncoderDestroy(m_dptr->enc_handle);
        m_dptr->enc_handle = nullptr;
    }
    if(m_dptr->output_buf){
        delete m_dptr->output_buf;
        m_dptr->output_buf = nullptr;
    }

    delete m_dptr;
}

bool JP2_U16C1Encoder::isValid()
{
    return m_dptr->isvalid;
}

bool JP2_U16C1Encoder::setLossy(bool islossy)
{
    if(!m_dptr->isvalid) return false;
    m_dptr->enc_config.irreversible = islossy ? 1 : 0;
    auto j2k_ret = nvjpeg2kEncodeParamsSetEncodeConfig(m_dptr->enc_params,
                                                       &m_dptr->enc_config);
    return NVJPEG2K_STATUS_SUCCESS == j2k_ret;
}

bool JP2_U16C1Encoder::doEncode(const void *input_buf, double psnr)
{
    if(!input_buf) return false;
    if(!m_dptr->isvalid) return false;

    if(psnr>0){
        m_dptr->psnr = psnr;
    }else{
        m_dptr->psnr = 0;
    }

    auto j2k_ret = nvjpeg2kEncodeParamsSetQuality(m_dptr->enc_params,m_dptr->psnr);
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    auto cuda_ret = cudaMemcpy2DAsync(m_dptr->input_image.pixel_data[0],
            m_dptr->input_image.pitch_in_bytes[0],
            input_buf, 2 * m_dptr->image_comp_info.component_width ,
            m_dptr->image_comp_info.component_width * 2,
            m_dptr->image_comp_info.component_height,
            cudaMemcpyHostToDevice,m_dptr->cudastream
            );
    if(cuda_ret!=cudaSuccess) return false;

    j2k_ret = nvjpeg2kEncode(m_dptr->enc_handle,
                             m_dptr->enc_state,
                             m_dptr->enc_params,
                             &m_dptr->input_image,
                             m_dptr->cudastream); //encode
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    // set the compressed_data buffer to NULL to retrieve the bitstream size
    j2k_ret = nvjpeg2kEncodeRetrieveBitstream(m_dptr->enc_handle,
                                              m_dptr->enc_state,
                                              nullptr,
                                              &m_dptr->encoded_size,
                                              m_dptr->cudastream);
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    if(m_dptr->encoded_size > m_dptr->output_buf_size){
        if(m_dptr->output_buf){
            delete m_dptr->output_buf;
            m_dptr->output_buf = nullptr;
        }
        m_dptr->output_buf_size = m_dptr->encoded_size;
        m_dptr->output_buf = new unsigned char[m_dptr->output_buf_size];
    }

    j2k_ret = nvjpeg2kEncodeRetrieveBitstream(m_dptr->enc_handle,
                                              m_dptr->enc_state,
                                              m_dptr->output_buf,
                                              &m_dptr->encoded_size,
                                              m_dptr->cudastream);
    if(j2k_ret!=NVJPEG2K_STATUS_SUCCESS) return false;

    cuda_ret = cudaStreamSynchronize(m_dptr->cudastream);
    if(cuda_ret!=cudaSuccess) return false;

    return true;
}

const char *JP2_U16C1Encoder::getOutputBuffer(size_t &size)
{
    if(!m_dptr->isvalid){
        size = 0;
        return nullptr;
    }

    size = m_dptr->encoded_size;
    return reinterpret_cast<const char *>(m_dptr->output_buf);
}
